import pandas as pd

from models.openai_client import generate_sql
from data_processor.excel_handler import ExcelHandler
from utils.sql_validator import is_valid_sql


class NL2SQLPipeline:
    def __init__(self):
        self.excel_handler = ExcelHandler()

    def run(self, nl_query: str) -> tuple[pd.DataFrame | None, dict]:
        info = {}   # 记录过程信息

        # 生成SQL
        sql = generate_sql(nl_query)
        info["sql"] = sql

        # 校验SQL
        is_valid, valid_msg = is_valid_sql(sql)
        info["validation"] = {"is_valid": is_valid, "valid_msg": valid_msg}
        if not is_valid:
            return None, info

        # 执行SQL
        result_df, exec_msg = self.excel_handler.execute_sql(sql)
        info["execution"] = exec_msg

        return result_df, info


if __name__ == "__main__":
    pipeline = NL2SQLPipeline()
    test_query = "查2018-9-18东北的前10个的订单信息"
    result, info = pipeline.run(test_query)
    print("过程信息：", info)
    print("最终结果：", result)
